import asyncio
import fileinput
import os

import pandas
import tqdm

from api import chat, init_models, init_index
from configs.model_config import EMBEDDING_MODEL

init_models(".vscode/rocket_qa")
# init_index("prd", ".vscode/rocket_qa/OBM_TEST")
output_file = ".vscode/test.xlsx"
if os.path.exists(output_file):
    dataframe = pandas.read_excel(output_file)
else:
    dataframe = pandas.DataFrame()

query = []
content = []
for line in tqdm.tqdm(fileinput.input("../langchain-ChatGLM/.vscode/input.txt")):
    result = asyncio.run(chat("OBM_TEST", line.strip(), dense_only=False))
    query.append(line.strip())
    content.append("\n".join([i.raw_content for i in result.source_documents])),


dataframe["query"] = query
dataframe[EMBEDDING_MODEL] = content
dataframe.to_excel(".vscode/test.xlsx", index=False)
